# Copyright 2017 Neural Networks and Deep Learning lab, MIPT
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, softwaredata
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import copy
import json
from logging import getLogger
from pathlib import Path
from typing import Dict, List

log = getLogger(__name__)


class SimpleDSTC2DatasetReader:
    @staticmethod
    def _data_fname(datatype):
        assert datatype in ('trn', 'val', 'tst'), "wrong datatype name"
        return f"simple-dstc2-{datatype}.json"

    def read(self, data_path: str, dialogs: bool = False) -> Dict[str, List]:
        """
        Parameters:
            data_path: DSTC2 dataset
            dialogs: flag which indicates whether to output list of turns or
             list of dialogs

        Returns:
            dictionary of dialogs from train/valid/test
            Each field is a list of tuples ``(user turn, system turn)``.
        """
        data = {
            'train': self._read_from_file(Path(data_path, self._data_fname('trn')), dialogs),
            'valid': self._read_from_file(Path(data_path, self._data_fname('val')), dialogs),
            'test':  self._read_from_file(Path(data_path, self._data_fname('tst')), dialogs)
        }
        log.info(f"There are {len(data['train'])} samples in train split.")
        log.info(f"There are {len(data['valid'])} samples in valid split.")
        log.info(f"There are {len(data['test'])} samples in test split.")
        return data

    @classmethod
    def _read_from_file(cls, file_path: str, dialogs: bool = False):
        """
            Returns data from single file
            dialogs:
                True 时返回对话, [[(q,a), (q,a), ...], ...]
                False 时返回 qa, [(q,a), (q,a), ...]
        """
        log.info(f"[loading dialogs from {file_path}]")

        utterances, responses, dialog_indices =\
            cls._get_turns(json.load(open(file_path, 'rt')), with_indices=True)

        data = list(map(cls._format_turn, zip(utterances, responses)))

        if dialogs:
            return [data[idx['start']:idx['end']] for idx in dialog_indices]
        return data

    @staticmethod
    def _format_turn(turn):
        turn_x, turn_y = turn
        x = {'text': turn_x['text']}
        y = {'text': turn_y['text'],
             'act': turn_y['act']}
        if 'act' in turn_x:
            x['intents'] = turn_x['act']    # 动作即意图
        if 'episode_done' in turn_x:
            x['episode_done'] = turn_x['episode_done']
        if turn_x.get('db_result') is not None:
            x['db_result'] = turn_x['db_result']
        if turn_x.get('slots'):
            x['slots'] = turn_x['slots']
        if turn_y.get('slots'):
            y['slots'] = turn_y['slots']
        return (x, y)

    @staticmethod
    def _get_turns(data, with_indices=False):
        n = 0
        utterances, responses, dialog_indices = [], [], []
        for dialog in data:
            cur_n_utter, cur_n_resp = 0, 0
            for i, turn in enumerate(dialog):
                speaker = turn.pop('speaker')
                if speaker == 1:    # 客户
                    if i == 0:
                        turn['episode_done'] = True  # 一通对话，由客户开始进行对话是不正常的，首先应当是机器人的欢迎语
                    utterances.append(turn)
                    cur_n_utter += 1    # 当前的对话轮次
                elif speaker == 2:  # 机器人方
                    responses.append(turn)
                    cur_n_resp += 1     # 系统回复轮次
                    if cur_n_utter not in range(cur_n_resp - 2, cur_n_resp + 1):    # 基本上算是一人一句，机器方由于要查询系统，可以存在连续动作,
                        raise RuntimeError("Datafile has wrong format.")            # 另外，由于系统初始是机器先说话，故　cur_n_utter <= cur_n_resp
                    if cur_n_utter != cur_n_resp:   # 二者不相等时， 构建虚拟客户对话，始终保持，human--bot 一一对应
                        if i == 0:  # 机器欢迎语，在这先 human 没有说话
                            new_utter = {
                                "text": "",
                                "episode_done": True
                            }
                        else:  # 机器查询数据库，导致 bot 连续两次对话
                            new_utter = copy.deepcopy(utterances[-1])
                            if 'db_result' not in responses[-2]:
                                raise RuntimeError("Every api_call action"
                                                   " should have db_result")
                            db_result = responses[-2].pop('db_result')
                            new_utter['db_result'] = db_result
                        utterances.append(new_utter)
                        cur_n_utter += 1        # 避免出现多次操作数据库导致的　utter 和 resp　相差过大
            if cur_n_utter != cur_n_resp:
                raise RuntimeError("Datafile has wrong format.")
            n += cur_n_utter
            dialog_indices.append({
                'start': n - cur_n_utter,
                'end': n,
            })

        if with_indices:
            return utterances, responses, dialog_indices
        return utterances, responses


if __name__ == '__main__':
    data = SimpleDSTC2DatasetReader().read('../../../my_data')
